#include "TRestTask.h"
#include "TRestRun.h"
#include "TRestG4Metadata.h"
#include "TRestG4Event.h"
#include "TRestHitsEvent.h"
#include <fstream>
#include <sstream>
#include "TFile.h"
#include "TTree.h"
#include <TRandom3.h>
#include <TMath.h>
#include <cmath>
#include "toolkitOfRest.C"

const float SRM_XY = 1;//mm
const float SRM_Z = 2;


// 计算距离
float calculateDistance(float x1, float y1, float z1, float x2, float y2, float z2) 
{
    float dx = x2 - x1;
    float dy = y2 - y1;
    float dz = z2 - z1;
    return std::sqrt(dx*dx + dy*dy + dz*dz);
}

int combineGasAndCZT()
{
    TTree* eventTree = getEventTree("/lustre/pandax/wangxuan/restTest/MeGat/v3_ManaRoot/Mana_e_z-_1MeV_n1w_1mmPixel.root",
                                    "EventTree");
    string outRootName = "/lustre/pandax/wangxuan/restTest/MeGat/v2_cgc/GasCZTCombine_e_z-_1MeV_n1w.root";
    float CUT_ENERGY = 0;

    TRestG4Event* g4Event = new TRestG4Event();
    TRestHitsEvent* hitsEvent = new TRestHitsEvent();
    eventTree->SetBranchAddress("TRestG4EventBranch", & g4Event);
    eventTree->SetBranchAddress("TRestHitsEventBranch", & hitsEvent);

    TRestG4Track* g4Track = new TRestG4Track();
    TRestG4Hits* g4Hits = new TRestG4Hits();

    int num_entries = eventTree->GetEntries();

    int eventID;
    float x;
    float y;
    float z;
    float e;
    vector<float> fx;
    vector<float> fy;
    vector<float> fz;
    vector<float> fe;

    vector<float> gas_x;
    vector<float> gas_y;
    vector<float> gas_z;
    vector<float> gas_e;
    float x_min;
    float y_min;
    float z_min;
    float x_max;
    float y_max;
    float z_max;
    float dx;
    float dy;
    float dz;
    vector<float> czt_x;
    vector<float> czt_y;
    vector<float> czt_z;
    vector<float> czt_e;
    
    float gamma_x;
    float gamma_y;
    float gamma_z;

    vector<float> G4_electron_track_x;
    vector<float> G4_electron_track_y;
    vector<float> G4_electron_track_z;
    vector<float> G4_electron_track_e;

    float G4_electron_energy;
    float electron_energy;
    float gamma_energy;
    float gasDepositedEnergy;
    float cztDepositedEnergy;
    float totalDepositedEnergy;
    float totalDepositedEnergy_smear;
    float fTotalDepositedEnergy;
    float fTotalDepositedEnergy_smear;
    TVector3 priDirection;
    float G4_e_vertex_x = 0;
    float G4_e_vertex_y = 0;
    float G4_e_vertex_z = 0;

    float boundary_center_x = 0;
    float boundary_center_y = 0;
    float boundary_center_z = 0;

    float boundary_e = 0;

    float boundary_distance = 0;
    vector<float> vector_boundry_distance;
    const float CUT_DISTANCE = 100;
    
    TFile * outRootfile = new TFile(outRootName.c_str() , "RECREATE");
    TTree * tree  = new TTree( "CombineHits" , "Hits Position Data" );
    tree->Branch("id", & eventID);

    // 将Gas与CZT中的点放在一个vector中 *没有采取这种做法
    // tree->Branch("fx", & fx);
    // tree->Branch("fy", & fy);
    // tree->Branch("fz", & fz);
    // tree->Branch("fe", & fe);

    // 将Gas与CZT中的点分开存放
    tree->Branch("gas_x", & gas_x);
    tree->Branch("gas_y", & gas_y);
    tree->Branch("gas_z", & gas_z);
    tree->Branch("gas_e", & gas_e);
    // 气体部分中径迹的纵深
    tree->Branch("dx", & dx);
    tree->Branch("dy", & dy);
    tree->Branch("dz", & dz);

    tree->Branch("czt_x", & czt_x);
    tree->Branch("czt_y", & czt_y);
    tree->Branch("czt_z", & czt_z);
    tree->Branch("czt_e", & czt_e);

    // 记录restG4中寻得的电子方向
    tree->Branch("priDirection_x", & priDirection[0]);
    tree->Branch("priDirection_y", & priDirection[1]);
    tree->Branch("priDirection_z", & priDirection[2]);

    // 记录restG4中寻得的电子径迹
    tree->Branch("G4_electron_track_x", & G4_electron_track_x);
    tree->Branch("G4_electron_track_y", & G4_electron_track_y);
    tree->Branch("G4_electron_track_z", & G4_electron_track_z);
    tree->Branch("G4_electron_track_e", & G4_electron_track_e);

    // 记录restG4中寻得的顶点，后续做为重建顶点的参照
    tree->Branch("G4_e_vertex_x", & G4_e_vertex_x);
    tree->Branch("G4_e_vertex_y", & G4_e_vertex_y);
    tree->Branch("G4_e_vertex_z", & G4_e_vertex_z);

    // 记录gamma沉积点的位置
    tree->Branch("gamma_x", & gamma_x);
    tree->Branch("gamma_y", & gamma_y);
    tree->Branch("gamma_z", & gamma_z);

    // 记录边界点
    tree->Branch("boundary_center_x", & boundary_center_x);
    tree->Branch("boundary_center_y", & boundary_center_y);
    tree->Branch("boundary_center_z", & boundary_center_z);

    
    // restG4电子径迹中进入到CZT中的点到边界点的距离
    tree->Branch("vector_boundry_distance", & vector_boundry_distance);

    tree->Branch("G4_electron_energy", & G4_electron_energy); // G4中找到的电子能量
    tree->Branch("electron_energy", & electron_energy);       // 修正后的反冲电子能量
    tree->Branch("gamma_energy", & gamma_energy);             // 修正后的散射光子能量
    tree->Branch("gasDepositedEnergy", & gasDepositedEnergy); // 原模拟中气体的能量
    tree->Branch("cztDepositedEnergy", & cztDepositedEnergy); // 原模拟中CZT的能量
    tree->Branch("totalDepositedEnergy", & totalDepositedEnergy); // Gas+CZT的能量总沉积
    tree->Branch("fTotalDepositedEnergy", & fTotalDepositedEnergy); // 原模拟的能量总沉积
    tree->Branch("totalDepositedEnergy_smear", & totalDepositedEnergy_smear); // 加入能量分辨后的能量总沉积
    tree->Branch("fTotalDepositedEnergy_smear", & fTotalDepositedEnergy_smear); 

    for (int i=0; i<10; i++)
    {
        // init

        // fx.clear();
        // fy.clear();
        // fz.clear();
        // fe.clear();
        bool FIND_ELECTRON = false;
        bool FIND_CZT = false;
        bool FIND_GAMMA = false;

        gas_x.clear();
        gas_y.clear();
        gas_z.clear();
        gas_e.clear();

        czt_x.clear();
        czt_y.clear();
        czt_z.clear();
        czt_e.clear();
        
        G4_electron_track_x.clear();
        G4_electron_track_y.clear();
        G4_electron_track_z.clear();
        G4_electron_track_e.clear();

        G4_e_vertex_x = 0;
        G4_e_vertex_y = 0;
        G4_e_vertex_z = 0;

        electron_energy = 0;
        gamma_energy = 0;
        gasDepositedEnergy = 0;
        cztDepositedEnergy = 0;
        totalDepositedEnergy = 0;
        totalDepositedEnergy_smear = 0;
        G4_electron_energy = 0;

        boundary_center_x = 0;
        boundary_center_y = 0;
        boundary_center_z = 0;
        boundary_e = 0;
    
        bool EXIST_BOUNDARY_POINT = false;
        vector_boundry_distance.clear();

        // 流程从这里开始
        eventTree->GetEntry(i);
        eventID = i;

        int num_of_tracks = g4Event->GetNumberOfTracks();
        fTotalDepositedEnergy = g4Event->GetTotalDepositedEnergy();
        priDirection = g4Event->GetPrimaryEventDirection(0);
        fTotalDepositedEnergy_smear = smearing(fTotalDepositedEnergy,1e3,1.5);
        // 遍历restManager branch中的点，这一部分是只有气体中的扩散电子
        int num_of_mana_hits = hitsEvent->GetNumberOfHits();
        float totalCharge = 0;
        x_min = hitsEvent->GetX(0);
        y_min = hitsEvent->GetY(0);
        z_min = hitsEvent->GetZ(0);
        x_max = hitsEvent->GetX(0);
        y_max = hitsEvent->GetY(0);
        z_max = hitsEvent->GetZ(0); 
        for (int j=0; j<num_of_mana_hits; j++)
        {   
            e = hitsEvent->GetEnergy(j);
            if (e > 0)
            {
                x = hitsEvent->GetX(j);
                y = hitsEvent->GetY(j);
                z = hitsEvent->GetZ(j);
                
                // 计算Gas与CZT交界中的点的电荷中心
                if(abs(x)>149 || abs(y)>149 || z<-149)
                {
                    // cout  << "Entry_" << i << ", boundary point: " << "(" << x << ", " << y << ", " << z << ")" << endl;
                    boundary_center_x += x*e;
                    boundary_center_y += y*e;
                    boundary_center_z += z*e;
                    boundary_e += e;
                    EXIST_BOUNDARY_POINT = true;
                }

                // 记录Gas中的所有点，添加能量smear
                if (abs(x)<=150 && abs(y)<=150 && abs(z)<=150)
                {
                    totalDepositedEnergy += e;
                    e = smearing(e, 1e3, 1.5);
                    totalDepositedEnergy_smear += e;
                    gasDepositedEnergy += e;
                    electron_energy += e;
                    // fx.push_back(x); 
                    // fy.push_back(y); 
                    // fz.push_back(z); 
                    // fe.push_back(e);

                    if (x < x_min) x_min = x;
                    else if (x > x_max) x_max = x;
                    if (y < y_min) y_min = y;
                    else if (y > y_max) y_max = y;
                    if (z < z_min) z_min = z;
                    else if (z > z_max) z_max = z;
                    gas_x.push_back(x); 
                    gas_y.push_back(y); 
                    gas_z.push_back(z); 
                    gas_e.push_back(e);

                }   
            }           
        }
        dx = x_max - x_min;
        dy = y_max - y_min;
        dz = z_max - z_min;

        if (EXIST_BOUNDARY_POINT)
        {
            boundary_center_x /= boundary_e;
            boundary_center_y /= boundary_e;
            boundary_center_z /= boundary_e;
        }
        
        // hit点最多的径迹认为是电子主径迹，遍历它及后续的次生径迹
        int electron_track_id = getMainTrackID(g4Event);
        for (int j=electron_track_id; j<num_of_tracks-1; j++)
        {
            g4Track = g4Event->GetTrack(j);
            g4Hits = g4Track->GetHits();
            G4_electron_energy += smearing(getTrackEnergy(g4Track), 1e3, 1.5);
            int num_of_hits = g4Track->GetNumberOfHits();
            if (num_of_hits != 0)
            {
                for (int j=0; j<num_of_hits; j++)
                {
                    e = g4Hits->GetEnergy(j);
                    if (e > 0)
                    {
                        x = g4Hits->GetX(j);
                        y = g4Hits->GetY(j);
                        z = g4Hits->GetZ(j);
                        G4_electron_track_x.push_back(x);
                        G4_electron_track_y.push_back(y);
                        G4_electron_track_z.push_back(z);
                        G4_electron_track_e.push_back(e);
                    }

                    // 计算restG4中电子径迹进入CZT中的点到边界点的距离
                    if (EXIST_BOUNDARY_POINT && (z<-150 || abs(x)>150 || abs(y)>150))
                    {
                        boundary_distance = sqrt(pow((x - boundary_center_x), 2) + pow((y - boundary_center_y), 2) + pow((z - boundary_center_z), 2));
                        vector_boundry_distance.push_back(boundary_distance);
                    }
                    
                }
            }
            
        }

        // 遍历restG4 branch中的点，将坐标位于czt部分的进行记录
        for (int j=0; j<num_of_tracks; j++)
        {   
            g4Track = g4Event->GetTrack(j);
            g4Hits = g4Track->GetHits();
            int num_of_hits = g4Track->GetNumberOfHits();
            float tmp_distance = 500;
            for (int k=0; k<num_of_hits; k++)
            {   
                e = g4Hits->GetEnergy(k);
                if (e > 0)
                {
                    x = g4Hits->GetX(k);
                    y = g4Hits->GetY(k);
                    z = g4Hits->GetZ(k);

                    // restG4散射电子的位置
                    if (j==0 && k==0)
                    {
                        G4_e_vertex_x = x;
                        G4_e_vertex_y = y;
                        G4_e_vertex_z = z;
                        FIND_ELECTRON = true;
                    }

                    // z<-150就是在底部的CZT中
                    if (z < -150)
                    {   
                        // 栅格化+能量Smear
                        totalDepositedEnergy += e;
                        x = ceilToInt((x-SRM_XY/2)/SRM_XY) * SRM_XY;
                        y = ceilToInt((y-SRM_XY/2)/SRM_XY) * SRM_XY;
                        z = ceilToInt((z-SRM_Z/2)/SRM_Z) * SRM_Z;
                        e = smearing(e, 1e3, 1.5);
                        totalDepositedEnergy_smear += e;
                        cztDepositedEnergy += e;

                        // fx.push_back(x); 
                        // fy.push_back(y); 
                        // fz.push_back(z); 
                        // fe.push_back(e); 

                        czt_x.push_back(x); 
                        czt_y.push_back(y); 
                        czt_z.push_back(z); 
                        czt_e.push_back(e);
                        FIND_CZT = true;
                    }
                    // 如果x或y的坐标abs大于150，那么就是在侧边CZT中。由于是竖着的，栅格化要变
                    else if (abs(x) > 150)
                    {
                        totalDepositedEnergy += e;
                        x = ceilToInt((x-SRM_XY/2)/SRM_XY) * SRM_XY;
                        y = ceilToInt((y-SRM_Z/2)/SRM_Z) * SRM_Z;
                        z = ceilToInt((z-SRM_XY/2)/SRM_XY) * SRM_XY;
                        e = smearing(e, 1e3, 1.5);
                        totalDepositedEnergy_smear += e;
                        cztDepositedEnergy += e;

                        // fx.push_back(x); 
                        // fy.push_back(y); 
                        // fz.push_back(z); 
                        // fe.push_back(e);

                        czt_x.push_back(x); 
                        czt_y.push_back(y); 
                        czt_z.push_back(z); 
                        czt_e.push_back(e);
                        FIND_CZT = true;

                    }
                    else if (abs(y) > 150)
                    {
                        totalDepositedEnergy += e;
                        x = ceilToInt((x-SRM_Z/2)/SRM_Z) * SRM_Z;
                        y = ceilToInt((y-SRM_XY/2)/SRM_XY) * SRM_XY;
                        z = ceilToInt((z-SRM_XY/2)/SRM_XY) * SRM_XY;
                        e = smearing(e, 1e3, 1.5);
                        totalDepositedEnergy_smear += e;
                        cztDepositedEnergy += e;

                        // fx.push_back(x); 
                        // fy.push_back(y); 
                        // fz.push_back(z); 
                        // fe.push_back(e);

                        czt_x.push_back(x); 
                        czt_y.push_back(y); 
                        czt_z.push_back(z); 
                        czt_e.push_back(e);
                        FIND_CZT = true;
                    }
                }          
            }
        }

        if (totalDepositedEnergy_smear > CUT_ENERGY)
        {   
            tree->Fill();
        }   
        
        if (i%500==0)
        {
            cout << "processing event_" << i << endl;
        }
    }
    tree->Write();
    outRootfile->Close();
    return 0;
}